'''
https://leetcode.cn/problems/count-submatrices-with-all-ones/description/
'''
from collections import deque
from typing import List


class Solution:
    def numSubmat(self, mat: List[List[int]]) -> int:
        def count(histogram):
            # 统计直方图📊中，必须以最后一行结尾，共有多少个子矩阵
            n = len(histogram)
            stack = deque()     # 找两侧比我小的元素。 大压小的栈
            ans = 0
            for i, height in enumerate(histogram):
                while stack and height < histogram[stack[-1]]:
                    cur = stack.pop()
                    left = stack[-1] if stack else -1
                    right = i
                    # 比如我是 9 ,左右两侧比我小的分别是 l位置的2, r位置7,
                    # 那我负责统计 l+1 ... r-1 范围内 8,9高度的答案
                    count_height_range = histogram[cur] - max(histogram[right], 0 if left == -1 else histogram[left])
                    count_range = right - left - 1
                    # 有多少个呢，共有是 (r-l-1)个1个宽度的 (r-l-2)个2个宽度的 (r-l-3) 个3个宽度的。。。
                    #               即 1, 2, 3, ... , count_range
                    #           然后这是一个高度的，共count_height_range个高度需要这样统计的
                    ans += count_height_range * (1 + count_range) / 2 * count_range
                stack.append(i)
            while stack:
                cur = stack.pop()
                left = stack[-1] if stack else -1
                right = n
                # 同样的统计逻辑
                count_height_range = histogram[cur] - max(0 if right == n else histogram[right], 0 if left == -1 else histogram[left])
                count_range = right - left - 1
                ans += count_height_range * (1 + count_range) / 2 * count_range
            return ans

        m, n = len(mat), len(mat[0])
        histogram = [0] * n
        ans = 0
        for i in range(m):
            for j in range(n):
                if mat[i][j] == 0:
                    histogram[j] = 0
                else:
                    histogram[j] += 1
            print(histogram)
            ans += count(histogram)
        return int(ans)